from nlu.pipe.extractors.extractor_base_data_classes import SparkNLPExtractor, SparkNLPExtractorConfig
from nlu.pipe.extractors.extractor_methods.helper_extractor_methods import *
from nlu.pipe.extractors.extractor_methods.base_extractor_methods import *

"""
This file contains methods to get pre-defined configurations for every annotator in OCR
Extractor_resolver.py should be used to resolve SparkNLP Annotator classes to methods 
in this file, which return the corrosponding configs that need to be passed to 
the master_extractor() call.

This file is where all the in extractor_base_data_classes.py Dataclasses are combined with the 
extractors defined in extractor_methods.py.

"""
def default_text_recognizer_config(output_col_prefix='ocr_text'):
    return SparkOCRExtractorConfig(
        get_text=True,
        output_col_prefix=output_col_prefix,
        name='default text recognizer config',
        description='Gets the text from an text recognizing OCR transformer',
    )


def default_visual_classifier_config(output_col_prefix='visual_classifier'):
    return SparkOCRExtractorConfig(
        output_col_prefix=output_col_prefix,
        name='visual_classifier label and confidence ',
        description='Gets label and confidence of visual classifier',
    )

def default_form_relation_extractor_config(output_col_prefix='extracted_relations'):
    return SparkOCRExtractorConfig(
        output_col_prefix=output_col_prefix,
        get_result=True,
        get_full_meta=True,
        name='full_relation_extraction',
        description='Get relation extraction result and all metadata, with positions of entities',
    )


def default_visual_ner_config(output_col_prefix='visual_ocr'):
    return SparkOCRExtractorConfig(
        get_text=True,
        get_begin=True,
        get_end=True,
        get_result=True,
        get_meta=True,
        get_full_meta=True,
        get_image_data=True,
        get_path=True,
        get_annotator_type=False,
        output_col_prefix=output_col_prefix,
        meta_white_list=['entity', 'confidence', 'sentence', 'chunk'],
        name='visual_ner label, confidence and entities ',
        description='Gets label, entities and confidence of visual ner',
    )

def default_binary_to_image_config(output_col_prefix='binary_image'):
    return SparkOCRExtractorConfig(
        output_col_prefix=output_col_prefix,
        get_image_origin=True,
        get_image_height=True,
        get_image_width=True,
        get_image_n_channels=True,
        get_image_mode=True,
        get_image_resolution=True,
        get_image_data = True,
        name='FULL binary to image extractor ',
        description='Gets all fields generated by binary to image Transformer ',
    )

def default_pdf_to_image_config(output_col_prefix='pdf_image'):
    return SparkOCRExtractorConfig(
        output_col_prefix=output_col_prefix,
        get_image_origin=True,
        get_image_height=True,
        get_image_width=True,
        get_image_n_channels=True,
        get_image_mode=True,
        get_image_resolution=True,
        get_image_data = True,
        get_img_positions = True,
        name='FULL pdf to image extractor ',
        description='Gets all fields generated by pdf to image Transformer ',
    )

def default_position_finder_config(output_col_prefix='positions'):
    return SparkOCRExtractorConfig(
        output_col_prefix=output_col_prefix,
        get_image_origin=True,
        get_image_height=True,
        get_image_width=True,
        get_image_n_channels=True,
        get_image_mode=True,
        get_image_resolution=True,
        get_image_data = True,
        get_img_positions = True,
        name='FULL binary to image extractor ',
        description='Gets all fields generated by binary to image Transformer ',
    )